import pandas
from sklearn import tree
import pydotplus
from sklearn.tree import DecisionTreeClassifier
import matplotlib.pyplot as plt
import matplotlib.image as pltimg
import os     
#引入Graphviz绘制决策树图片的程序的地址
os.environ["PATH"] += os.pathsep + 'D:/Graphviz/bin'

#如需制作决策树，所有数据都必须是数字。

df = pandas.read_csv("shows.csv")

#将字符串更改为数值
d={'UK':0,'USA':1,'N':2}
df['Nationality']=df['Nationality'].map(d)
d={'YES':1,'NO':0}
df['Go']=df['Go'].map(d)

#分开特征列和目标列
features=['Age','Experience','Rank','Nationality']

X=df[features]
y=df['Go']

print(X)
print(y)

#创建决策树
dtree=DecisionTreeClassifier()
dtree=dtree.fit(X,y)
data=tree.export_graphviz(dtree,out_file=None,feature_names=features)
graph=pydotplus.graph_from_dot_data(data)
graph.write_png("mydecisiontree.png")

img=pltimg.imread("mydecisiontree.png")
imgplot=plt.imshow(img)
plt.show()
